# @Time : 2021/8/4 19:14
# @Author : Li Kunlun
# @Description : 模型构造

from mxnet import nd
from mxnet.gluon import nn


# 1、继承Block类来构造模型
class MLP(nn.Block):
    # 声明带有模型参数的层，这里声明了两个全连接层
    def __init__(self, **kwargs):
        # 调用MLP父类Block的构造函数来进行必要的初始化。这样在构造实例时还可以指定其他函数
        super(MLP, self).__init__(**kwargs)
        self.hidden = nn.Dense(256, activation='relu')  # 隐藏层
        self.output = nn.Dense(10)  # 输出层

    # 定义模型的前向计算，即如何根据输入x计算返回所需要的模型输出
    def forward(self, x):
        return self.output(self.hidden(x))


X = nd.random.uniform(shape=(2, 20))
net = MLP()
net.initialize()
"""
[[ 0.09543003  0.04614331 -0.00286653 -0.07790348 -0.0513024   0.02942039
   0.08696644 -0.01907929 -0.04122178  0.05088577]
 [ 0.07692869  0.03099705  0.00856576 -0.044672   -0.06926841  0.09132433
   0.06786595 -0.06187843 -0.03436673  0.04234697]]
<NDArray 2x10 @cpu(0)>
"""
print(net(X))


# 2、Sequential类继承自Block类
class MySequential(nn.Block):
    def __init__(self, **kwargs):
        super(MySequential, self).__init__(**kwargs)

    def add(self, block):
        """
        :param block: block是一个Block子类实例，假设它有一个独一无二的名字。我们将它保存在Block类的
            成员变量_children里，其类型是OrderedDict。当MySequential实例调用
            initialize函数时，系统会自动对_children里所有成员初始化
        :return:
        """
        self._children[block.name] = block

    def forward(self, x):
        # OrderedDict保证会按照成员添加时的顺序遍历成员
        for block in self._children.values():
            x = block(x)
        return x


net = MySequential()
net.add(nn.Dense(256, activation='relu'))
net.add(nn.Dense(10))
net.initialize()
"""
[[ 0.00362229  0.00633331  0.03201145 -0.01369376  0.10336448 -0.03508019
  -0.00032164 -0.01676024  0.06978627  0.01303309]
 [ 0.03871719  0.02608211  0.03544958 -0.02521311  0.11005434 -0.01430662
  -0.03052466 -0.03852825  0.06321152  0.00385941]]
<NDArray 2x10 @cpu(0)>
"""
print(net(X))


# 3、构造复杂的模型
class FancyMLP(nn.Block):
    def __init__(self, **kwargs):
        super(FancyMLP, self).__init__(**kwargs)
        # 使用get_constant创建的随机权重参数不会在训练中被迭代（即常数参数）
        self.rand_weight = self.params.get_constant('rand_weight', nd.random.uniform(shape=(20, 20)))
        self.dense = nn.Dense(20, activation='relu')

    def forward(self, x):
        x = self.dense(x)
        # 使用创建的常数参数，以及NDArray的relu函数和dot函数
        x = nd.relu(nd.dot(x, self.rand_weight.data()) + 1)
        # 复用全连接层。等价于两个全连接层共享参数
        x = self.dense(x)
        # 控制流，这里我们需要调用asscalar函数来返回标量进行比较
        while x.norm().asscalar() > 1:
            x /= 2
        if x.norm().asscalar() < 0.8:
            x *= 10
        return x.sum()


net = FancyMLP()
net.initialize()

# [ 18.57195282]
# <NDArray 1 @cpu(0)>
print(net(X))
